package tools;
import java.util.*;

public class AssemblyCodeGenerator {
    private final Set<String> variables = new LinkedHashSet<>();
    private final Set<String> labels = new HashSet<>();
    private final StringBuilder dataSegment = new StringBuilder();
    private final StringBuilder codeSegment = new StringBuilder();
    private boolean hasPrint = false;
    private int tempCounter = 0;

    public String generate(List<String> intermediateCode) {
        initializeDataSegment();
        analyzeCode(intermediateCode);
        processIntermediateCode(intermediateCode);
        return buildFinalAssemblyCode();
    }

    private void initializeDataSegment() {
        dataSegment.append(".model small\n.stack 100h\n");
        dataSegment.append(".data\n");
    }

    private void analyzeCode(List<String> intermediateCode) {
        for (String line : intermediateCode) {
            line = line.trim();
            if (line.isEmpty()) continue;

            // 处理变量声明 int a
            if (line.matches("int\\s+\\w+")) {
                String var = line.split("\\s+")[1];
                variables.add(var);
            }
            // 处理带初始化的声明 int b = 5
            else if (line.matches("int\\s+\\w+\\s*=\\s*\\d+")) {
                String[] parts = line.split("\\s*=\\s*");
                String var = parts[0].split("\\s+")[1];
                variables.add(var);
            }
            else if (line.matches("[a-zA-Z]\\w*\\s*=\\s*[\\w\\d]+")) {
                String var = line.split("=")[0].trim();
                variables.add(var);
            }
            else if (line.matches("if\\s+.*")) {
                String[] parts = line.split("\\s+");
                variables.add(parts[1]);
                if (!parts[3].matches("\\d+")) {
                    variables.add(parts[3]);
                }
            }
            else if (line.matches("\\w+:")) {
                labels.add(line.replace(":", ""));
            }
            else if (line.startsWith("printf")) {
                hasPrint = true;
                variables.add(line.substring(6).trim());
            }
            else if (line.matches("t\\d+\\s*=\\s*.*")) {
                String tempVar = line.split("=")[0].trim();
                if (!tempVar.matches("t\\d+")) {
                    throw new RuntimeException("Invalid temporary variable: " + tempVar);
                }
                int num = Integer.parseInt(tempVar.substring(1));
                if (num > tempCounter) {
                    tempCounter = num;
                }
            }
        }
    }

    private void processIntermediateCode(List<String> intermediateCode) {
        for (String line : intermediateCode) {
            line = line.trim();
            if (line.isEmpty()) continue;

            if (line.matches("int\\s+\\w+")) {
                // 简单声明 int a，不需要生成代码
                continue;
            }
            else if (line.matches("int\\s+\\w+\\s*=\\s*\\d+")) {
                handleInitializedDeclaration(line);
            }
            else if (line.matches("[a-zA-Z]\\w*\\s*=\\s*[\\w\\d]+")) {
                handleVariableAssignment(line);
            }
            else if (line.startsWith("printf")) {
                handlePrintStatement(line);
            }
            else if (line.matches("\\w+:")) {
                codeSegment.append(line).append("\n");
            }
            else if (line.startsWith("if")) {
                handleConditionalJump(line);
            }
            else if (line.startsWith("goto")) {
                handleUnconditionalJump(line);
            }
            else if (line.matches("t\\d+\\s*=\\s*.*")) {
                handleTempAssignment(line);
            }
        }
    }

    private void handleInitializedDeclaration(String line) {
        String[] parts = line.split("\\s*=\\s*");
        String varPart = parts[0].trim();
        String var = varPart.split("\\s+")[1];
        String value = parts[1].trim();

        dataSegment.append(String.format("%s dw %s\n", var, value));
    }

    private void handleVariableAssignment(String line) {
        String[] parts = line.split("=");
        String var = parts[0].trim();
        String value = parts[1].trim();

        if (value.matches("\\d+")) {
            codeSegment.append(String.format("    mov ax, %s\n", value));
        } else {
            codeSegment.append(String.format("    mov ax, [%s]\n", value));
        }
        codeSegment.append(String.format("    mov [%s], ax\n", var));
    }

    private void handlePrintStatement(String line) {
        hasPrint = true;
        String var = line.substring(6).trim();

        codeSegment.append(String.format("    ; Print %s\n", var));
        codeSegment.append(String.format("    mov ax, [%s]\n", var));
        codeSegment.append("    call print_number\n");
    }

    private void handleConditionalJump(String line) {
        String[] parts = line.split("\\s+");
        String var = parts[1];
        String operator = parts[2];
        String operand = parts[3];
        String label = parts[5];

        codeSegment.append(String.format("    mov ax, [%s]\n", var));

        if (operand.matches("\\d+")) {
            codeSegment.append(String.format("    cmp ax, %s\n", operand));
        } else {
            codeSegment.append(String.format("    mov bx, [%s]\n", operand));
            codeSegment.append("    cmp ax, bx\n");
        }

        switch (operator) {
            case "<": codeSegment.append(String.format("    jl %s\n", label)); break;
            case "<=": codeSegment.append(String.format("    jle %s\n", label)); break;
            case ">": codeSegment.append(String.format("    jg %s\n", label)); break;
            case ">=": codeSegment.append(String.format("    jge %s\n", label)); break;
            case "==": codeSegment.append(String.format("    je %s\n", label)); break;
            case "!=": codeSegment.append(String.format("    jne %s\n", label)); break;
        }
    }

    private void handleUnconditionalJump(String line) {
        String label = line.split("\\s+")[1];
        codeSegment.append(String.format("    jmp %s\n", label));
    }

    private void handleTempAssignment(String line) {
        String[] parts = line.split("=");
        String tempVar = parts[0].trim();
        String expr = parts[1].trim();

        if (expr.matches("[\\w\\d]+\\s*[+\\-*/]\\s*[\\w\\d]+")) {
            String[] ops = expr.split("(?<=[+\\-*/])|(?=[+\\-*/])");
            String left = ops[0].trim();
            String op = ops[1].trim();
            String right = ops[2].trim();

            // 加载左操作数
            if (left.matches("\\d+")) {
                codeSegment.append(String.format("    mov ax, %s\n", left));
            } else {
                codeSegment.append(String.format("    mov ax, [%s]\n", left));
            }

            // 处理运算
            if (right.matches("\\d+")) {
                switch (op) {
                    case "+": codeSegment.append(String.format("    add ax, %s\n", right)); break;
                    case "-": codeSegment.append(String.format("    sub ax, %s\n", right)); break;
                    case "*":
                        codeSegment.append(String.format("    mov bx, %s\n", right));
                        codeSegment.append("    imul bx\n");
                        break;
                    case "/":
                        codeSegment.append("    cwd\n");
                        codeSegment.append(String.format("    mov bx, %s\n", right));
                        codeSegment.append("    idiv bx\n");
                        break;
                }
            } else {
                codeSegment.append(String.format("    mov bx, [%s]\n", right));
                switch (op) {
                    case "+": codeSegment.append("    add ax, bx\n"); break;
                    case "-": codeSegment.append("    sub ax, bx\n"); break;
                    case "*": codeSegment.append("    imul bx\n"); break;
                    case "/":
                        codeSegment.append("    cwd\n");
                        codeSegment.append("    idiv bx\n");
                        break;
                }
            }

            // 存储结果
            codeSegment.append(String.format("    mov [%s], ax\n", tempVar));
        }
    }

    private String buildFinalAssemblyCode() {
        StringBuilder finalCode = new StringBuilder();
        finalCode.append(dataSegment.toString());

        // 声明所有临时变量
        for (int i = 0; i <= tempCounter; i++) {
            finalCode.append(String.format("t%d dw ?\n", i));
        }

        // 声明所有已识别但未初始化的变量
        for (String var : variables) {
            if (!var.matches("t\\d+") && !dataSegment.toString().contains(var + " dw ")) {
                finalCode.append(String.format("%s dw ?\n", var));
            }
        }

        // 添加共用的打印资源
        if (hasPrint) {
            finalCode.append("print_msg db 'Value = $'\n");
            finalCode.append("print_buf db 6 dup('$')\n");
        }

        // 代码段
        finalCode.append(".code\n");
        finalCode.append("main proc\n");
        finalCode.append("    mov ax, @data\n");
        finalCode.append("    mov ds, ax\n");
        finalCode.append(codeSegment.toString());


        finalCode.append("\n    mov ax, 4C00h\n");
        finalCode.append("    int 21h\n");
        finalCode.append("main endp\n");

        // 添加共用的打印函数
        if (hasPrint) {
            finalCode.append("\n; ===== 共用打印函数 =====\n");
            finalCode.append("print_number proc\n");
            finalCode.append("    push ax\n");
            finalCode.append("    push bx\n");
            finalCode.append("    push cx\n");
            finalCode.append("    push dx\n");
            finalCode.append("    push di\n");
            finalCode.append("    push si\n");


            // 打印前缀
//            finalCode.append("    lea dx, print_msg\n");
//            finalCode.append("    mov ah, 09h\n");
//            finalCode.append("    int 21h\n");

            // 数字转字符串
            finalCode.append("    mov cx, 10\n");
            finalCode.append("    lea di, print_buf+5\n");
            finalCode.append("    mov byte ptr [di], '$'\n");
            finalCode.append("    mov si,ax\n");


            // 处理负数
            finalCode.append("    test ax, ax\n");
            finalCode.append("    jns positive\n");
            finalCode.append("    neg ax\n");
            finalCode.append("    push ax\n");
            finalCode.append("    mov dl, '-'\n");
            finalCode.append("    mov ah, 02h\n");
            finalCode.append("    int 21h\n");
            finalCode.append("    pop ax\n");
            finalCode.append("positive:\n");

            // 转换循环
            finalCode.append("convert_loop:\n");
            finalCode.append("    xor dx, dx\n");
            finalCode.append("    div cx\n");
            finalCode.append("    add dl, '0'\n");
            finalCode.append("    dec di\n");
            finalCode.append("    mov [di], dl\n");
            finalCode.append("    test ax, ax\n");
            finalCode.append("    jnz convert_loop\n");

            // 打印结果
            finalCode.append("    mov dx, di\n");
            finalCode.append("    mov ah, 09h\n");
            finalCode.append("    int 21h\n");

            // 换行
            finalCode.append("    mov dl, 13\n");
            finalCode.append("    mov ah, 02h\n");
            finalCode.append("    int 21h\n");
            finalCode.append("    mov dl, 10\n");
            finalCode.append("    int 21h\n");

            finalCode.append("    pop si\n");
            finalCode.append("    pop di\n");
            finalCode.append("    pop dx\n");
            finalCode.append("    pop cx\n");
            finalCode.append("    pop bx\n");
            finalCode.append("    pop ax\n");
            finalCode.append("    ret\n");
            finalCode.append("print_number endp\n");
        }


        finalCode.append("end main\n");

        return finalCode.toString();
    }
}